/**
 * Copyright (c) 2023 murenchao
 * taomu is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *       http://license.coscl.org.cn/MulanPubL-2.0
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PubL v2 for more details.
 */
package cool.taomu.toolkit.crypto.util

import cool.taomu.toolkit.crypto.Base64
import java.io.BufferedWriter
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.FileWriter
import java.io.OutputStream
import java.nio.charset.Charset
import java.security.SecureRandom
import java.security.Security
import java.util.UUID
import org.apache.commons.io.FileUtils
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.openpgp.PGPCompressedData
import org.bouncycastle.openpgp.PGPCompressedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedData
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedDataList
import org.bouncycastle.openpgp.PGPLiteralData
import org.bouncycastle.openpgp.PGPOnePassSignatureList
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection
import org.bouncycastle.openpgp.PGPSignatureList
import org.bouncycastle.openpgp.PGPUtil
import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory
import org.bouncycastle.openpgp.operator.jcajce.JcaKeyFingerprintCalculator
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyDecryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePGPDataEncryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyDataDecryptorFactoryBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyKeyEncryptionMethodGenerator
import org.bouncycastle.util.io.Streams
import org.slf4j.LoggerFactory

class PgpUtils {
	val static LOG = LoggerFactory.getLogger(PgpUtils);

	def static void main(String[] args) {
		Security.addProvider(new BouncyCastleProvider());
		var String pk = FileUtils.readFileToString(new File("src/test/resources/publicKey1.txt"),
			Charset.forName("UTF-8"));
		var data = PgpUtils.encrypt("Hello World!!!! ".bytes, pk, true);
		println(new String(data, "UTF-8"));
		var String prk = FileUtils.readFileToString(new File("src/test/resources/privateKey1.txt"),
			Charset.forName("UTF-8"));
		var d = PgpUtils.decrypt(data, prk, "123456789");
		println("result : " + new String(d, "UTF-8"));
	}

	def static decrypt(byte[] data, String privateKey, String password) {

		try(var bais = new ByteArrayInputStream(data)) {
			var input = PGPUtil.getDecoderStream(bais);
			var jpof = new JcaPGPObjectFactory(input);
			var o = jpof.nextObject;
			var PGPEncryptedDataList enc;
			if (o instanceof PGPEncryptedDataList) {
				enc = o;
			} else {
				enc = jpof.nextObject as PGPEncryptedDataList;
			}
			// 获取privatekey
			var pkeyBase64 = new Base64(privateKey.bytes).decode();
			val pgpSec = new PGPSecretKeyRingCollection(pkeyBase64, new JcaKeyFingerprintCalculator());
			var edo = enc.encryptedDataObjects
			var pbe = edo.findFirst [ k |
				pgpSec.getSecretKey((k as PGPPublicKeyEncryptedData).keyID) !== null
			] as PGPPublicKeyEncryptedData;
			var secKey = pgpSec.getSecretKey(pbe.keyID);
			var jpekdb = new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray);
			var sKey = secKey.extractPrivateKey(jpekdb);
			var clear = pbe.getDataStream(new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(sKey));
			var plain = new JcaPGPObjectFactory(clear);
			var message = plain.nextObject
			if (input !== null) {
				input.close();
			}
			try(var out = new ByteArrayOutputStream()) {
				if (message instanceof PGPCompressedData) {
					var fact = new JcaPGPObjectFactory(message.dataStream);
					message = fact.nextObject;
				}
				if (message instanceof PGPLiteralData) {
					Streams.pipeAll(message.inputStream, out);
				} else if (message instanceof PGPOnePassSignatureList) {
					LOG.info("encrypted message contains a signed message - not literal data.")
				} else if (message instanceof PGPSignatureList) {
					LOG.info("encrypted message contains a signed message - not literal data.");
				}
				if (pbe.isIntegrityProtected) {
					if (!pbe.verify) {
						LOG.info("message failed integrity check");
					} else {
						LOG.info("message integrity check passed");
					}
				} else {
					LOG.info("no message integrity check");
				}
				return out.toByteArray
			}
		}
	}

	def static encrypt(byte[] data, String publicKey, boolean withIntegrityCheck) {
		try(var out = new ByteArrayOutputStream()) {
			encrypt(data, publicKey, withIntegrityCheck, out);
			return out.toByteArray;
		}
	}

	def static encrypt(byte[] data, String publicKey, boolean withIntegrityCheck, OutputStream out) {
		// 获取publiekey
		try(var bais = new ByteArrayInputStream(new Base64(publicKey.bytes).decode)) {
			var input = PGPUtil.getDecoderStream(bais);
			var pkrc = new PGPPublicKeyRingCollection(input, new JcaKeyFingerprintCalculator);
			var encKey = pkrc.keyRings.findFirst [
				it.publicKeys.findFirst[it.isEncryptionKey].isEncryptionKey;
			].publicKey;
			var file = CreateTempFile.create(data, "UTF-8");
			try(var baos = new ByteArrayOutputStream()) {
				var cdg = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
				PGPUtil.writeFileToLiteralData(cdg.open(baos), 'b', file);
				cdg.close;
				var bytes = baos.toByteArray;
				var jpdeb = new JcePGPDataEncryptorBuilder(PGPEncryptedData.CAST5);
				jpdeb = jpdeb.setWithIntegrityPacket(withIntegrityCheck);
				var bc = jpdeb.setSecureRandom(new SecureRandom()).setProvider("BC");
				var encGen = new PGPEncryptedDataGenerator(bc);
				encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(encKey).provider = "BC");
				try(var cout = encGenOpen(encGen, out, bytes)) {
					cout.write(bytes);
				}
			} finally {
				file.delete
			}
		}
	}

	private def static OutputStream encGenOpen(PGPEncryptedDataGenerator encGen, OutputStream out, byte[] bytes) {
		encGen.open(out, bytes.length)
	}

	static class CreateTempFile {
		def static create(byte[] content, String charset) {
			return create(new String(content, Charset.forName(charset)));
		}

		def static create(String content) {
			var tempName = UUID.randomUUID.toString();
			var temp = File.createTempFile(tempName, ".temp");
			temp.deleteOnExit();
			try(var bw = new BufferedWriter(new FileWriter(temp))) {
				bw.write(content);
			}
			return temp;
		}

		def static create(byte[] content) {
			var tempName = UUID.randomUUID.toString();
			var temp = File.createTempFile(tempName, ".temp");
			temp.deleteOnExit();
			FileUtils.writeByteArrayToFile(temp, content);
			return temp;
		}
	}
}
